# -*- coding: utf-8 -*-
from datasets import load_dataset

class SFTDataProcessor:
    """
    数据处理器类，负责数据集的加载、格式化和Token化。
    """

    def __init__(self, tokenizer, max_seq_length=1024):
        """
        初始化SFTDataProcessor

        Args:
            tokenizer: 用于文本tokenization的tokenizer
            max_seq_length (int): 序列的最大长度，默认为256
        """
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length

    def _format_prompt(self, example, include_response=True):
        """
        格式化单个样本的prompt

        Args:
            example (dict): 包含instruction、input和output键的字典
            include_response (bool): 是否在prompt末尾包含响应内容

        Returns:
            str: 格式化后的prompt字符串
        """
        # 检查是否有input字段且不为空
        if 'input' in example and example['input']:
            prompt = (
                f"Below is an instruction that describes a task, paired with an input that provides further context. "
                f"Write a response that appropriately completes the request.\n\n"
                f"### Instruction:\n{example['instruction']}\n\n"
                f"### Input:\n{example['input']}\n\n"
                f"### Response:\n"
            )
        else:
            prompt = (
                f"Below is an instruction that describes a task. "
                f"Write a response that appropriately completes the request.\n\n"
                f"### Instruction:\n{example['instruction']}\n\n"
                f"### Response:\n"
            )

        if include_response:
            prompt += example['output']

        return prompt

    def _tokenize_function(self, examples):
        """
        对训练样本进行tokenization处理

        Args:
            examples (dict): 包含样本数据的字典

        Returns:
            dict: 包含input_ids、attention_mask和labels的字典
        """
        # 格式化prompt并在末尾添加结束符
        full_prompt = self._format_prompt(examples, include_response=True) + self.tokenizer.eos_token

        # 使用tokenizer处理文本
        result = self.tokenizer(
            full_prompt,
            truncation=True,
            max_length=self.max_seq_length,
            padding="max_length"
        )

        # 将input_ids复制为labels用于训练
        result["labels"] = result["input_ids"].copy()
        return result

    def load_and_prepare_dataset(self, file_path, num_samples=None):
        """
        从本地JSONL文件加载数据集并进行预处理。

        Args:
            file_path (str): 本地JSONL文件的路径。
            num_samples (int, optional): 要使用的训练样本数量。如果为None，则使用整个数据集。

        Returns:
            Dataset: 处理后的数据集
        """
        print(f"正在从本地文件加载数据集: {file_path}")
        dataset = load_dataset('json', data_files=file_path, split="train")

        if num_samples and num_samples < len(dataset):
            print(f"使用 {num_samples} 条数据。")
            dataset = dataset.select(range(num_samples))
        else:
            print(f"使用全部 {len(dataset)} 条数据。")

        print("正在对数据集进行Tokenize...")
        tokenized_dataset = dataset.map(self._tokenize_function, remove_columns=list(dataset.features))
        return tokenized_dataset